#!/usr/bin/env python3
# SPDX-License-Identifier: None

import os
import re
import networkx as nx
import matplotlib.pyplot as plt
from typing import Dict, List, Set, Tuple, Optional
import argparse
import sys

# 导入现有的解析器类
from zenmpx.event_parser import SapphireRapidsEventParser, SapphireRapidsMetricParser


class EventTreeNode:
    """事件树中的节点，表示一个事件"""

    def __init__(self, event_name: str, event_data: Dict = None):
        self.name = event_name
        self.data = event_data or {}
        self.connections: List[Tuple[EventTreeNode, float]] = []

    def add_connection(self, node: "EventTreeNode", weight: float = 1.0):
        """添加到另一个节点的连接"""
        if node != self and not any(
            conn[0].name == node.name for conn in self.connections
        ):
            self.connections.append((node, weight))
            # 也在目标节点上添加反向连接
            node.add_connection(self, weight)

    def to_dict(self) -> Dict:
        """将节点转换为字典格式"""
        return {
            "name": self.name,
            "data": self.data,
            "connections": [conn[0].name for conn in self.connections],
        }


class EventTree:
    """事件树类，用于构建和管理事件之间的关系"""

    def __init__(self):
        self.nodes: Dict[str, EventTreeNode] = {}
        self.graph = nx.Graph()

    def add_event(self, event_name: str, event_data: Dict = None):
        """添加一个事件节点"""
        if event_name not in self.nodes:
            self.nodes[event_name] = EventTreeNode(event_name, event_data)
            self.graph.add_node(event_name, **(event_data or {}))
        return self.nodes[event_name]

    def get_event(self, event_name: str) -> Optional[EventTreeNode]:
        """获取指定名称的事件节点"""
        return self.nodes.get(event_name)

    def add_connection(self, event_name1: str, event_name2: str, weight: float = 1.0):
        """添加两个事件之间的连接"""
        if event_name1 not in self.nodes:
            self.add_event(event_name1)
        if event_name2 not in self.nodes:
            self.add_event(event_name2)

        # 在节点对象上添加连接
        node1 = self.nodes[event_name1]
        node2 = self.nodes[event_name2]
        node1.add_connection(node2, weight)

        # 在networkx图上添加边
        self.graph.add_edge(event_name1, event_name2, weight=weight)

    def build_from_events_and_metrics(
        self, events_data: List[Dict], metrics_data: List[Dict]
    ):
        """从事件和度量数据构建事件树"""
        # 1. 首先添加所有事件作为节点
        for event in events_data:
            event_name = event.get("EventName", str(event.get("EventCode", "unknown")))
            self.add_event(event_name, event)

        # 2. 处理度量数据，为每个度量中引用的事件添加连接
        for metric in metrics_data:
            metric_name = metric.get("MetricName", "unknown_metric")
            metric_expr = metric.get("MetricExpr", "")

            # 从表达式中提取事件名称
            event_names = self._extract_event_names_from_expr(metric_expr)

            # 为该度量中的所有事件对添加连接
            for i in range(len(event_names)):
                for j in range(i + 1, len(event_names)):
                    # 可以选择增加权重，例如基于度量的重要性或出现频率
                    self.add_connection(event_names[i], event_names[j], weight=1.0)

    def _extract_event_names_from_expr(self, expr: str) -> List[str]:
        """从度量表达式中提取事件名称"""
        # 使用正则表达式尝试匹配典型的事件名称格式
        event_patterns = [
            # 匹配类似INST_RETIRED.ANY这样的事件名称
            r"\b[A-Z][A-Z0-9_]+(\.[A-Z][A-Z0-9_]+)+\b",
            # 匹配类似CPU_CLK_UNHALTED.THREAD的事件名称
            r"\b[A-Z]+_[A-Z]+(?:_[A-Z]+)*\.[A-Z]+(?:_[A-Z]+)*\b",
        ]

        event_names = []
        for pattern in event_patterns:
            matches = re.findall(pattern, expr)
            if isinstance(matches, list):
                event_names.extend(matches)
            else:
                event_names.append(matches)

        # 去重并返回
        return list(set(event_names))

    def visualize(self, output_file: str = None):
        """可视化事件树"""
        plt.figure(figsize=(12, 10))

        # 使用spring布局算法
        pos = nx.spring_layout(self.graph, k=0.3)

        # 绘制节点和边
        nx.draw_networkx_nodes(self.graph, pos, node_size=300, node_color="skyblue")
        nx.draw_networkx_edges(self.graph, pos, alpha=0.5)

        # 添加标签，但只为部分节点添加以避免拥挤
        if len(self.graph.nodes) <= 50:
            nx.draw_networkx_labels(self.graph, pos, font_size=8)
        else:
            # 只显示一些关键节点的标签
            labels = {
                node: node if len(node) < 15 else node[:12] + "..."
                for node in self.graph.nodes
            }
            nx.draw_networkx_labels(self.graph, pos, labels, font_size=6)

        plt.title("Events Relationship Tree")
        plt.axis("off")

        if output_file:
            plt.savefig(output_file, dpi=300, bbox_inches="tight")
            print(f"可视化结果已保存到: {output_file}")
        else:
            plt.show()


class EventTreeBuilder:
    """事件树构建器，用于从文件中加载数据并构建事件树"""

    def __init__(self):
        self.event_tree = EventTree()

    def load_events_data(self, events_dir: str) -> List[Dict]:
        """使用现有的事件解析器加载事件数据"""
        all_events = []

        # 加载所有事件文件
        for filename in os.listdir(events_dir):
            if filename.endswith(".json"):
                file_path = os.path.join(events_dir, filename)
                try:
                    # 使用现有的SapphireRapidsEventParser来解析事件文件
                    parser = SapphireRapidsEventParser(file_path)
                    all_events.extend(parser.events)
                except Exception as e:
                    print(f"解析文件 {file_path} 时出错: {e}")
                    continue

        return all_events

    def load_metrics_data(self, metrics_dir: str) -> List[Dict]:
        """使用现有的度量解析器加载度量数据"""
        all_metrics = []

        # 递归遍历所有子目录查找度量文件
        for root, _, files in os.walk(metrics_dir):
            for filename in files:
                if filename.endswith(".json"):
                    file_path = os.path.join(root, filename)
                    try:
                        # 使用现有的SapphireRapidsMetricParser来解析度量文件
                        parser = SapphireRapidsMetricParser(file_path)
                        all_metrics.extend(parser.metrics)
                    except Exception as e:
                        print(f"解析文件 {file_path} 时出错: {e}")
                        continue

        return all_metrics

    def build_from_directory(self, data_dir: str) -> EventTree:
        """从目录中构建事件树"""
        events_dir = os.path.join(data_dir, "events")
        metrics_dir = os.path.join(data_dir, "metrics")

        # 检查目录是否存在
        if not os.path.exists(events_dir):
            raise FileNotFoundError(f"事件目录不存在: {events_dir}")
        if not os.path.exists(metrics_dir):
            raise FileNotFoundError(f"度量目录不存在: {metrics_dir}")

        # 加载数据
        print(f"正在加载事件数据...")
        events_data = self.load_events_data(events_dir)
        print(f"已加载 {len(events_data)} 个事件")

        print(f"正在加载度量数据...")
        metrics_data = self.load_metrics_data(metrics_dir)
        print(f"已加载 {len(metrics_data)} 个度量")

        # 构建事件树
        print(f"正在构建事件树...")
        self.event_tree.build_from_events_and_metrics(events_data, metrics_data)

        return self.event_tree


class EnhancedEventTreeBuilder(EventTreeBuilder):
    """增强版事件树构建器，提供更高级的功能"""

    def __init__(self):
        super().__init__()

    def build_core_only_tree(self, data_dir: str) -> EventTree:
        """构建只包含核心事件的树"""
        events_dir = os.path.join(data_dir, "events")
        metrics_dir = os.path.join(data_dir, "metrics")

        # 加载所有数据
        all_events = []
        for filename in os.listdir(events_dir):
            if filename.endswith(".json") and "core" in filename.lower():
                file_path = os.path.join(events_dir, filename)
                try:
                    parser = SapphireRapidsEventParser(file_path)
                    all_events.extend(parser.events)
                except Exception as e:
                    print(f"解析文件 {file_path} 时出错: {e}")
                    continue

        # 加载所有度量数据
        all_metrics = self.load_metrics_data(metrics_dir)

        # 构建事件树
        self.event_tree.build_from_events_and_metrics(all_events, all_metrics)
        return self.event_tree

    def build_high_weight_connections(
        self, data_dir: str, weight_threshold: int = 3
    ) -> EventTree:
        """构建连接权重更高的事件树"""
        events_data = self.load_events_data(data_dir)
        metrics_data = self.load_metrics_data(data_dir)

        # 先构建基础树
        self.event_tree.build_from_events_and_metrics(events_data, metrics_data)

        # 计算连接频率
        connection_count = {}
        for metric in metrics_data:
            metric_expr = metric.get("MetricExpr", "")
            event_names = self.event_tree._extract_event_names_from_expr(metric_expr)

            # 计算每对事件的连接频率
            for i in range(len(event_names)):
                for j in range(i + 1, len(event_names)):
                    key = tuple(sorted([event_names[i], event_names[j]]))
                    connection_count[key] = connection_count.get(key, 0) + 1

        # 创建一个新的图，只保留高频连接
        self.event_tree = EventTree()

        # 先添加所有事件节点
        for event in events_data:
            event_name = event.get("EventName", str(event.get("EventCode", "unknown")))
            self.event_tree.add_event(event_name, event)

        # 只添加高频连接
        for (event1, event2), count in connection_count.items():
            if count >= weight_threshold:
                self.event_tree.add_connection(event1, event2, weight=count)

        return self.event_tree


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="构建和可视化事件关系树")
    parser.add_argument(
        "--data-dir",
        default="/home/xyjiang/project/zenmpx/data/spr/SPR",
        help="数据目录路径，默认为SPR数据",
    )
    parser.add_argument("--output", help="输出图片文件路径，不指定则直接显示")
    parser.add_argument(
        "--mode",
        choices=["normal", "core-only", "high-weight"],
        default="normal",
        help="构建模式：normal(普通模式), core-only(仅核心事件), high-weight(高频连接)",
    )
    parser.add_argument(
        "--weight-threshold",
        type=int,
        default=3,
        help="高频连接模式下的权重阈值，默认为3",
    )

    args = parser.parse_args()

    # 检查数据目录是否存在
    if not os.path.exists(args.data_dir):
        print(f"错误：数据目录 {args.data_dir} 不存在")
        sys.exit(1)

    print(f"正在从 {args.data_dir} 加载数据...")

    # 根据模式选择构建器
    if args.mode == "normal":
        # 构建事件树
        builder = EventTreeBuilder()
        event_tree = builder.build_from_directory(args.data_dir)
    elif args.mode == "core-only":
        # 构建仅包含核心事件的树
        builder = EnhancedEventTreeBuilder()
        event_tree = builder.build_core_only_tree(args.data_dir)
    else:  # high-weight
        # 构建高频连接的树
        builder = EnhancedEventTreeBuilder()
        event_tree = builder.build_high_weight_connections(
            args.data_dir, args.weight_threshold
        )

    print(
        f"事件树构建完成，包含 {len(event_tree.nodes)} 个节点和 {len(event_tree.graph.edges)} 条边"
    )

    # 可视化
    event_tree.visualize(args.output)


if __name__ == "__main__":
    main()
